import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import time

from matplotlib import rcParams

# ============================ 1. 设置超参数和配置 ============================

# 训练配置
BATCH_SIZE = 128  # 批大小
NUM_EPOCHS = 20  # 训练轮数
LEARNING_RATE = 0.001  # 学习率
MODEL_SAVE_PATH = './results/cifar100_model.pth'  # 模型保存路径
PLOT_SAVE_PATH = './results/accuracy_curve.png'  # 准确率曲线图保存路径

# CIFAR100数据集的均值和标准差（针对RGB通道）
MEAN = (0.5071, 0.4865, 0.4409)
STD = (0.2673, 0.2564, 0.2762)

# ============================ 2. 数据加载与预处理 ============================

# 数据预处理：数据增强和归一化
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),  # 随机裁剪
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),  # 转换为Tensor
    transforms.Normalize(MEAN, STD)  # 归一化
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(MEAN, STD)
])

# 下载并加载训练集和测试集
trainset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)

testset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=0)


# ============================ 3. 定义模型 ============================

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 卷积层1
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        # 卷积层2
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        # 卷积层3
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        # 池化层：尺寸减半
        self.pool = nn.MaxPool2d(2, 2)
        # 全连接层
        self.fc = nn.Linear(128 * 4 * 4, 100)  # CIFAR100有100类

    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))
        x = self.pool(x)  # 输出尺寸：32 x 16 x 16
        x = torch.relu(self.bn2(self.conv2(x)))
        x = self.pool(x)  # 输出尺寸：64 x 8 x 8
        x = torch.relu(self.bn3(self.conv3(x)))
        x = self.pool(x)  # 输出尺寸：128 x 4 x 4
        x = x.view(x.size(0), -1)  # 展开为一维
        x = self.fc(x)
        return x


# ============================ 4. 模型、损失函数与优化器设置 ============================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleCNN().to(device)

# 使用交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# Adam优化器
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)


# ============================ 5. 训练和测试过程 ============================

def train_one_epoch(epoch, model, trainloader, optimizer, criterion):
    model.train()
    running_loss = 0.0
    correct_train = 0
    total_train = 0
    for inputs, labels in trainloader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()  # 梯度清零
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

        running_loss += loss.item() * inputs.size(0)
        _, predicted = torch.max(outputs, 1)
        total_train += labels.size(0)
        correct_train += (predicted == labels).sum().item()

    train_acc = correct_train / total_train
    return running_loss / total_train, train_acc


def evaluate(model, testloader):
    model.eval()
    correct_test = 0
    total_test = 0
    with torch.no_grad():
        for inputs, labels in testloader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total_test += labels.size(0)
            correct_test += (predicted == labels).sum().item()

    test_acc = correct_test / total_test
    return test_acc


# ============================ 6. 训练与测试循环 ============================

train_accuracy_list = []
test_accuracy_list = []

# 创建保存结果的文件夹
if not os.path.exists('./results'):
    os.makedirs('./results')

# 训练和评估
for epoch in range(NUM_EPOCHS):
    # 训练一个epoch
    train_loss, train_acc = train_one_epoch(epoch, model, trainloader, optimizer, criterion)

    # 测试模型
    test_acc = evaluate(model, testloader)

    # 记录准确率
    train_accuracy_list.append(train_acc)
    test_accuracy_list.append(test_acc)

    print(f"Epoch [{epoch + 1}/{NUM_EPOCHS}], Loss: {train_loss:.4f}, "
          f"Train Accuracy: {train_acc:.4f}, Test Accuracy: {test_acc:.4f}")

# ============================ 7. 保存模型 ============================

# 保存训练后的模型
torch.save(model.state_dict(), MODEL_SAVE_PATH)
print(f"模型已保存至 {MODEL_SAVE_PATH}")

# ============================ 8. 绘制并保存准确率变化曲线 ============================

# 绘制准确率变化图

# 设置Matplotlib字体支持中文
rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 设置字体
rcParams['axes.unicode_minus'] = False  # 正常显示负号

# 绘制准确率变化图
plt.figure(figsize=(8, 6))
plt.plot(range(1, NUM_EPOCHS+1), train_accuracy_list, label='训练准确率')
plt.plot(range(1, NUM_EPOCHS+1), test_accuracy_list, label='测试准确率')
plt.xlabel('训练轮次')
plt.ylabel('准确率')
plt.title('CIFAR100 图像识别准确率变化')
plt.legend()
plt.grid(True)

# 保存准确率变化图
plt.savefig(PLOT_SAVE_PATH)
plt.show()
print(f"准确率变化图已保存至 {PLOT_SAVE_PATH}")
